#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : 5.PyTorch卷积层原理和使用.py
# @Author: Richard Chiming Xu
# @Date  : 2021/11/8
# @Desc  :

import torch
from torch import nn
from torchsummary import summary


class MyConvNet(torch.nn.Module):
    def __init__(self):
        super(MyConvNet, self).__init__()
        self.conv = nn.Conv2d(in_channels=1,out_channels=32,kernel_size=5,stride=1,padding=2)

    # forward 定义前向传播
    def forward(self, X):
        return self.conv(X)

model = MyConvNet()
model = model.cuda()


print(summary(model, input_size=(1, 32, 32), batch_size=-1))
print('参数量公式：out_channels*(in_channels*kernel_size_h*kernel_size_w)+bias={}*({}*{}*{})+{}={}'.format(32,1,5,5,32,32*(1*5*5)+32))

